# coding: utf-8

import datetime
import json
import time

import pandas as pd
import requests

from enums import CodeKind
from exts import ts, sql_pool, trade_cal


class Market(object):
    trade_cal = trade_cal
    instances = {}
    flag = {}

    def __new__(cls, code: str, name: str, start_date: str, end_date: str, *args, **kwargs):
        if code not in cls.instances:
            cls.instances[code] = super().__new__(cls)
            cls.flag[code] = False
        if len(cls.instances) > 100:  # 只保存最后使用的100只代码
            pop_code = list(cls.instances.keys())[0]
            cls.instances.pop(pop_code)
            cls.flag.pop(pop_code)
        return cls.instances[code]

    def __init__(self, code: str, name: str, start_date: str, end_date: str,
                 kind: CodeKind = CodeKind.STOCK, *args, **kwargs):
        """

        :param code: 代码
        :param name: 名称
        :param start_date: 开始日期 如20220101
        :param end_date: 结束日期 如20220131
        :param args:
        :param kwargs:
        """
        if self.flag[code]:  # 已经初始化
            return

        if start_date > end_date:
            start_date, end_date = end_date, start_date
        self.code = code
        self.kind = kind  # 类型： 股票  ETF  指数  基金
        self.code1 = self.change_code(code)  # 股票原基础上增加.SH  .SZ
        self.name = name
        self.start_date = None
        self.end_date = None
        self.market_data = pd.DataFrame()
        self._market_data(start_date, end_date)

        self.flag[code] = True

    def _market_data(self, start_date: str, end_date: str):
        """

        :param start_date: 格式如20220105
        :param end_date:
        :return:
        """
        if start_date > end_date:
            start_date, end_date = end_date, start_date

        if self.start_date is None or self.end_date is None:  # 初始化，start_date和end_date为None
            start = start_date
            end = end_date
        elif self.start_date <= start_date and self.end_date >= end_date:  # 数据已经读取到内存
            return
        else:
            start = min(self.start_date, start_date)
            end = max(self.end_date, end_date)
        if self.kind == CodeKind.STOCK:
            self.market_data = ts.daily(ts_code=self.code1, start_date=start, end_date=end)
            count = 3
            while self.market_data.empty and count > 0:
                time.sleep(1)
                count -= 1
                self.market_data = ts.daily(ts_code=self.code1, start_date=start, end_date=end)
        elif self.kind == CodeKind.INDEX:
            self.market_data = self.scrapy_index_market()
        elif self.kind == CodeKind.FUND:
            self.market_data = self.scrapy_fund_worth(start_date=start_date, end_date=end_date)
        try:
            self.market_data.sort_values(by="trade_date", inplace=True)
            self.start_date = self.market_data.iloc[0, :]["trade_date"]
            self.end_date = self.market_data.iloc[-1, :]["trade_date"]
        except Exception as e:
            print(self.code1, self.kind, "未查询到数据")
            print(str(e))

    def query_from_db(self, start_date: str, end_date: str):
        sql = """
                select * from fund_market where fund = %(code)s and trade_date >= %(start_date)s
                and trade_date <= %(end_date)s;
            """
        with sql_pool:
            sql_pool.execute(sql, {"code": self.code, "start_date": datetime.datetime.strptime(start_date, "%Y%m%d"),
                             "end_date": datetime.datetime.strptime(end_date, "%Y%m%d")})
            res = sql_pool.fetchall()
        df = pd.DataFrame(res)
        df["trade_date"] = df["trade_date"].apply(lambda x: x.strftime("%Y%m%d"))
        return df

    def scrapy_index_market(self):
        """
        爬取指数行情
        :return:
        """
        limit = 10000  # 从最新交易日往回去limit个交易日
        if self.code[0:3] == "000":  # 上证指数
            sec_id = "1." + self.code
        elif self.code[0:3] == "399":  # 深圳指数
            sec_id = "0." + self.code
        else:
            raise Exception("指数{index}暂不支持爬取行情".format(index=self.code))
        url = "http://83.push2his.eastmoney.com/api/qt/stock/kline/get?secid={sec_id}&fields1=f1,f2,f3,f4,f5,f6&" \
              "fields2=f51,f52,f53,f54,f55,f56,f57&klt=101&fqt=1&end=20500101&lmt={lmt}" \
            .format(sec_id=sec_id, lmt=limit)
        response = requests.get(url)
        res = json.loads(response.text)
        res = res["data"]
        self.name = res["name"]
        k_lines = res["klines"]
        k_lines = [dict(zip(["trade_date", "open", "close", "high", "low", "vol", "amount"], x.split(','))) for x in
                   k_lines]
        df = pd.DataFrame(k_lines)
        df["trade_date"] = df["trade_date"].str.replace("-", "")
        df[["open", "close", "high", "low", "vol", "amount"]] = \
            df[["open", "close", "high", "low", "vol", "amount"]].apply(pd.to_numeric, errors="coerce")
        return df

    def scrapy_fund_worth(self, start_date, end_date):
        """
        爬取基金净值
        :return:
        """
        total_page = 999
        page = 1
        page_size = 10000
        url = """http://api.fund.eastmoney.com/f10/lsjz"""
        header = {
            "Accept": "*/*",
            "Accept-Encoding": "gzip, deflate",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6",
            "Connection": "keep-alive",
            "Host": "api.fund.eastmoney.com",
            "Referer": "http://fundf10.eastmoney.com/"
        }
        df_return = pd.DataFrame()
        while page <= total_page:
            param = {
                "fundCode": self.code,
                "pageIndex": page,
                "pageSize": page_size,
                "startDate": datetime.datetime.strptime(start_date, "%Y%m%d").strftime("%Y-%m-%d"),
                "endDate": datetime.datetime.strptime(end_date, "%Y%m%d").strftime("%Y-%m-%d")
            }
            response = requests.get(url, params=param, headers=header)
            res = json.loads(response.text)

            total_page = int(res["TotalCount"] / page_size) + 1
            page += 1

            res = res["Data"]
            df = pd.DataFrame(res["LSJZList"])
            df_return = pd.concat([df_return, df])
        if df_return.empty:
            return df_return
        df_return.rename(columns={"FSRQ": "trade_date", "DWJZ": "unit_worth", "LJJZ": "cal_worth"}, inplace=True)
        df_return.drop_duplicates("trade_date", inplace=True)
        df_return.sort_values("trade_date", inplace=True)
        df_return["trade_date"] = df_return["trade_date"].str.replace("-", "")
        df_return[["unit_worth", "cal_worth"]] = \
            df_return[["unit_worth", "cal_worth"]].apply(pd.to_numeric, errors="coerce")
        return df_return

    def market_range_history(self, start_date: str, end_date: str):
        self._market_data(start_date, end_date)
        if self.market_data.empty:  # 做历史回测，可能etf基金还未成立
            return pd.DataFrame()
        return self.market_data[(self.market_data["trade_date"] >= start_date) &
                                (self.market_data["trade_date"] <= end_date)].copy(deep=True)

    def market_count_history(self, curr: str, count) -> pd.DataFrame:
        """
        获取前count个交易日历史行情
        :param curr: 当前日期
        :param count:
        :return:
        """
        end_date = datetime.datetime.strptime(curr, "%Y%m%d").strftime("%Y%m%d")
        start_date = self.trade_cal[(self.trade_cal["is_open"] == 1) &
                                    (self.trade_cal["cal_date"] <= end_date)][0:count]
        start_date = start_date.iloc[-1, :]["cal_date"]
        return self.market_range_history(start_date, end_date)

    def market_one_day(self, date: str):
        """
        获取某一天的行情
        :param date:
        :return:
        """
        df = self.market_range_history(start_date=date, end_date=date)
        if df.empty:
            return pd.Series(dtype="float64")
        else:
            return df.iloc[0, :]

    def change_code(self, code: str):
        """
        根据交易所地点，在原代码上添加后缀
        如上证交易所股票601108，添加.SH
        :param code: 代码
        :return:
        """
        if self.kind != CodeKind.STOCK:  # 只有股票，更换code1
            return code

        if ".SH" in code or ".SZ" in code:
            return code
        if code[0] == "6":  # 上海交易所
            return code + ".SH"
        elif code[0:2] == "00" or code[0:3] == "300":  # 深圳交易所
            return code + ".SZ"

    def next_trade_date(self, curr_date):
        trade_date = self.trade_cal[(self.trade_cal["is_open"] == 1) & (self.trade_cal["cal_date"] >= curr_date)]
        trade_date = trade_date.iloc[0, :]
        return trade_date["cal_date"]

    def count_trade_date(self, curr_date, num):
        """
        查找前num个交易日的日期
        :param curr_date
        :param num
        :return:
        """
        trade_date = self.trade_cal[(self.trade_cal["is_open"] == 1) & (self.trade_cal["cal_date"] <= curr_date)]
        trade_date = trade_date.iloc[num, :]
        return trade_date["cal_date"]


if __name__ == '__main__':
    a = Market("159985", "aaa", "20230101", "20240605", CodeKind.FUND)
    market = a.market_count_history("20240605", 100)
    print(market)
